
import { assert } from 'cc';
import { ITensorAllocator, ITensorData } from './Backends/ITensorAllocator';

/**
 * Max rank
 */
export const MaxRank = 8;

// The following dimension names are based on ONNX Dimension Denotation.
// see: https://github.com/onnx/onnx/blob/master/docs/DimensionDenotation.md

/**
 * Data channel dimension index number
 */
export const DataChannel = 7;
/**
 * Channels dimension index number
 */
export const C = DataChannel;

/**
 * Data feature 0 dimension index number
 */
export const DataFeature0 = 6;
/**
 * Width dimension index number
 */
export const W = DataFeature0;

/**
 * Data feature 1 dimension index number
 */
export const DataFeature1 = 5;
/**
 * Height dimension index number
 */
export const H = DataFeature1;

/**
 * Data feature 2 dimension index number
 */
export const DataFeature2 = 4;
/**
 * Depth dimension index number
 */
export const D = DataFeature2;

/**
 * Data feature 3 dimension index number
 */
export const DataFeature3 = 3;
/**
 * Batch dimension index number
 */
export const DataBatch = 2;

/**
 * Sequence length dimension index number
 */
export const NumberOfDirections = 1;

/**
 * Sequence length dimension index number
 */
export const SequenceLength = 0;

/**
 * Data features
 */
export const DataFeatures = [W, H, D, DataFeature3];

/**
 * Kernel input channel dimension
 */
export const KernelInChannel = 6;

/**
 * Kernel output channel dimension
 */
export const KernelOutChannel = 7;

/**
 * Kernel spatial dimension 0
 */
export const KernelSpatial0 = 5;

/**
 * Kernel spatial dimension 1
 */
export const KernelSpatial1 = DataBatch; // NOTE: maps to batch

/**
 * Kernel spatial dimension 2
 */
export const KernelSpatial2 = DataBatch - 1; // NOTE: maps to numDirections

/**
 * Kernel spatial dimension 3
 */
export const KernelSpatial3 = SequenceLength; // NOTE: maps to sequenceLength

/**
 * Kernel spatial dimensions
 */
export const KernelSpatials = [KernelSpatial0, KernelSpatial1, KernelSpatial2, KernelSpatial3];

/**
 * @en
 * Named dimension
 *
 * @zh
 * 命名维度
 */
export enum NamedDimension {
    S = 1 << SequenceLength,
    R = 1 << NumberOfDirections,
    N = 1 << DataBatch,
    T = 1 << DataFeature3,
    D = 1 << DataFeature2,
    H = 1 << DataFeature1,
    W = 1 << DataFeature0,
    C = 1 << DataChannel,

    None = 0,
    All = S | R | N | T | D | H | W | C,
}

/**
 * @en
 * Tensor shape
 * TensorShape are immutable representation of a Tensor dimensions and rank.
 * Depending on which constructor is used, the TensorShape will either be rank 8 and channels last (ie NHWC) or actual
 * rank with unnamed tensor dimensions when using the constructor that takes int[].
 *
 * @zh
 * 张量形状
 * TensorShape 是张量维度和秩的不可变表示。根据使用的构造函数，TensorShape 将是秩 8 和通道最后（即 NHWC）或实际秩，当使用 int[] 的构造函数时，张量维度是未命名的。
 *
 *
 */
export class TensorShape {
    /**
     * @en
     * The number of elements in the TensorShape
     *
     * @zh
     * 张量形状中的元素数量
     */
    private _data: number[] = [];
    public get data(): number[] {
        return this._data;
    }
    public set data(value: number[]) {
        this._data = value;
    }

    /**
     * @en
     * Return the number of sequence.
     * @zh
     * 返回序列的数量。
     */
    get sequenceLength(): number {
        if (this.hasNamedDimensions) {
            const value = this._data[SequenceLength];
            return value;
        }

        return 1;
    }

    /**
     * @en
     *  Return the number of direction.
     * @zh
     *  返回方向的数量。
     */
    get numberOfDirections(): number {
        if (this.hasNamedDimensions) {
            const value = this._data[NumberOfDirections];
            return value;
        }

        return 1;
    }

    /**
     * @en
     * Return the number of batch.
     * @zh
     * 返回批次的数量。
     */
    get batch(): number {
        if (this.hasNamedDimensions) {
            const value = this._data[DataBatch];
            return value;
        }

        return this.getAxis(0);
    }

    /**
     * @en
     * Return the size of 3rd spatial dimension (axis is DataFeature3)
     * @zh
     * 返回第三个空间维度的大小（轴是 DataFeature3）
     */
    get extraDimension(): number {
        if (this.hasNamedDimensions) {
            const value = this._data[DataFeature3];
            return value;
        }
        return 1;
    }

    /**
     * @en
     *  Return the spatial depth (axis is DataFeature2).
     * @zh
     * 返回空间深度（轴是 DataFeature2）。
     */
    get depth(): number {
        if (this.hasNamedDimensions) {
            const value = this._data[DataFeature2];
            return value;
        }
        return 1;
    }

    /**
     * @en
     *  Return the spatial height (axis is DataFeature1).
     * @zh
     * 返回空间高度（轴是 DataFeature1）。
     */
    get height(): number {
        if (this.hasNamedDimensions) {
            const value = this._data[DataFeature1];
            return value;
        }

        return this.getAxis(1);
    }

    /**
     * @en
     *  Return the spatial width (axis is DataFeature0).
     * @zh
     * 返回空间宽度（轴是 DataFeature0）。
     */
    get width(): number {
        if (this.hasNamedDimensions) {
            const value = this._data[DataFeature0];
            return value;
        }

        return this.getAxis(2);
    }

    /**
     * @en
     * Return the number of channels.
     * @zh
     * 返回通道数量
     */
    get channels(): number {
        if (this.hasNamedDimensions) {
            const value = this._data[DataChannel];
            return value;
        }

        return this.getAxis(3);
    }

    /**
     * @en
     * Always 8 if legacy, named constructors are used otherwise the actual rank.
     *Look also at the `dimensions` property.
     * @zh
     * 如果使用传统的命名构造函数，则始终为 8，否则为实际秩。还请查看 `dimensions` 属性。
     */
    private _rank: number = 0;
    get rank(): number {
        return this._rank;
    }
    set rank(value: number) {
        this._rank = value;
    }

    /**
     * @en
     * Whether this shape makes use of named dimensions or is nameless
     *
     * @zh
     * 此形状是否使用了命名维度或者是无名的
     */
    get hasNamedDimensions(): boolean {
        return this._usesNamedDimensions != 0;
    }

    private _usesNamedDimensions: NamedDimension = 0;
    public get usesNamedDimensions(): NamedDimension {
        return this._usesNamedDimensions;
    }
    public set usesNamedDimensions(value: NamedDimension) {
        this._usesNamedDimensions = value;
    }

    /**
     * @en
     * Kernel dimension ordering is [D,H,W,C,K] for efficiency purpose.
     * @zh
     * 为了效率目的，内核维度排序是 [D,H,W,C,K]。
     */
    get kernelSpatialDepth(): number {
        return this.numberOfDirections;
    }

    get kernelHeight(): number {
        return this.depth;
    }

    get kernelWidth(): number {
        return this.height;
    }

    get kernelDepth(): number {
        return this.width;
    }

    get kernelCount(): number {
        return this.channels;
    }

    get flatHeight(): number {
        return this.batch;
    }

    /**
     * @en
     *  Return the T*D*H*W*C.
     *
     * @zh
     * 返回 T*D*H*W*C。
     */
    get flatWidth(): number {
        let w = 1;
        if (this.hasNamedDimensions) {
            w = this.extraDimension * this.depth * this.height * this.width * this.channels;
            return w;
        }
        for (let i = 0; i < this.rank; i++) {
            w *= this.getAxis(i);
        }
        return w;
    }

    /**
     * @en
     *  Return the total number of elements represented by this shape.
     *
     * @zh
     * 返回此形状表示的元素的总数。
     */
    get length(): number {
        let l = 1;
        if (this.hasNamedDimensions) {
            l = this.sequenceLength * this.numberOfDirections * this.flatHeight * this.flatWidth;
            return l;
        }

        for (let i = 0; i < this.rank; i++) {
            l *= this.getAxis(i);
        }
        return l;
    }

    /**
     * @en
     * Return the count of non-unit dimension of this shape.
     * For example [N,1,1,C] dimensions is 2.
     *
     * @zh
     * 返回此形状的非单位维度计数。
     */
    get dimensions(): number {
        if (this.hasNamedDimensions) {
            return (
                (this.sequenceLength > 1 ? 1 : 0) +
                (this.numberOfDirections > 1 ? 1 : 0) +
                (this.batch > 1 ? 1 : 0) +
                (this.extraDimension > 1 ? 1 : 0) +
                (this.depth > 1 ? 1 : 0) +
                (this.height > 1 ? 1 : 0) +
                (this.width > 1 ? 1 : 0) +
                (this.channels > 1 ? 1 : 0)
            );
        }

        return this.rank;
    }

    static createWithArray(shape: number[], unnamedDimensions = false): TensorShape {
        assert(shape.length <= MaxRank, 'Shape length must be less than or equal to MaxRank 8');
        const tensorShape = new TensorShape();
        if (unnamedDimensions) {
            tensorShape.usesNamedDimensions = NamedDimension.None;
            tensorShape.rank = shape.length;

            if (tensorShape.rank > 0) {
                tensorShape.data = shape;
            } else {
                tensorShape.rank = 1;
                tensorShape.init(1, 1, 1, 1, 1, 1, 1, 1);
            }
        } else {
            switch (shape.length) {
                case 0:
                    tensorShape.usesNamedDimensions = NamedDimension.N;
                    tensorShape.init(1, 1, 1, 1, 1, 1, 1, 1);
                    break;
                case 1:
                    //[1,1,N,1,1,1,1,1]
                    tensorShape.usesNamedDimensions = NamedDimension.N;
                    tensorShape.init(1, 1, shape[0], 1, 1, 1, 1, 1);
                    break;
                case 2:
                    //[1,1,N,1,1,1,1,C]
                    tensorShape.usesNamedDimensions = NamedDimension.N | NamedDimension.C;
                    tensorShape.init(1, 1, 1, 1, 1, 1, shape[0], shape[1]);
                    break;
                case 3:
                    //[1,1,N,1,1,1,W,C]
                    tensorShape.usesNamedDimensions = NamedDimension.N | NamedDimension.W | NamedDimension.C;
                    tensorShape.init(1, 1, shape[0], 1, 1, 1, shape[1], shape[2]);
                    break;
                case 4:
                    //[1,1,N,1,1,H,W,C]
                    tensorShape.usesNamedDimensions =
                        NamedDimension.N | NamedDimension.H | NamedDimension.W | NamedDimension.C;
                    tensorShape.init(1, 1, shape[0], 1, 1, shape[1], shape[2], shape[3]);
                    break;
                case 5:
                    //[1,1,N,1,D,H,W,C]
                    tensorShape.usesNamedDimensions =
                        NamedDimension.N | NamedDimension.D | NamedDimension.H | NamedDimension.W | NamedDimension.C;
                    tensorShape.init(1, 1, shape[0], 1, shape[1], shape[2], shape[3], shape[4]);
                    break;
                case 6:
                case 7:
                    throw new Error(`Unsupported shape length ${shape.length}`);
                    break;
                case 8:
                    //[S,R,N,T,D,H,W,C]
                    tensorShape.usesNamedDimensions = NamedDimension.All;
                    tensorShape.init(
                        shape[SequenceLength],
                        shape[NumberOfDirections],
                        shape[DataBatch],
                        shape[DataFeature3],
                        shape[DataFeature2],
                        shape[DataFeature1],
                        shape[DataFeature0],
                        shape[DataChannel]
                    );
                    break;
            }
        }

        return tensorShape;
    }

    asNamed(): TensorShape {
        if (this.hasNamedDimensions) {
            return this;
        }

        const tensorShape = new TensorShape();
        switch (this.rank) {
            case 0:
                tensorShape.usesNamedDimensions = NamedDimension.N;
                tensorShape.init(1, 1, 1, 1, 1, 1, 1, 1);
                break;
            case 1:
                //[1,1,N,1,1,1,1,1]
                tensorShape.usesNamedDimensions = NamedDimension.N;
                tensorShape.init(1, 1, this.getAxis(0), 1, 1, 1, 1, 1);
                break;
            case 2:
                //[1,1,N,1,1,1,1,C]
                tensorShape.usesNamedDimensions = NamedDimension.N | NamedDimension.C;
                tensorShape.init(1, 1, 1, 1, 1, 1, this.getAxis(0), this.getAxis(1));
                break;
            case 3:
                //[1,1,N,1,1,1,W,C]
                tensorShape.usesNamedDimensions = NamedDimension.N | NamedDimension.W | NamedDimension.C;
                tensorShape.init(1, 1, this.getAxis(0), 1, 1, 1, this.getAxis(1), this.getAxis(2));
                break;
            case 4:
                //[1,1,N,1,1,H,W,C]
                tensorShape.usesNamedDimensions =
                    NamedDimension.N | NamedDimension.H | NamedDimension.W | NamedDimension.C;
                tensorShape.init(1, 1, this.getAxis(0), 1, 1, this.getAxis(1), this.getAxis(2), this.getAxis(3));
                break;
            case 5:
                //[1,1,N,1,D,H,W,C]
                tensorShape.usesNamedDimensions =
                    NamedDimension.N | NamedDimension.D | NamedDimension.H | NamedDimension.W | NamedDimension.C;
                tensorShape.init(
                    1,
                    1,
                    this.getAxis(0),
                    1,
                    this.getAxis(1),
                    this.getAxis(2),
                    this.getAxis(3),
                    this.getAxis(4)
                );
                break;
            case 6:
            case 7:
                throw new Error(`Unsupported shape length ${this.rank}`);
                break;
            case 8:
                //[S,R,N,T,D,H,W,C]
                tensorShape.usesNamedDimensions = NamedDimension.All;
                tensorShape.init(
                    this.getAxis(0),
                    this.getAxis(1),
                    this.getAxis(2),
                    this.getAxis(3),
                    this.getAxis(4),
                    this.getAxis(5),
                    this.getAxis(6),
                    this.getAxis(7)
                );
                break;
        }
        return tensorShape;
    }

    /**
     * @en
     *  Create a TensorShape of shape [S,R,N,T,D,H,W,C].
     * Currently seqLen must be 1.
     * @zh
     * 创建一个形状为 [S,R,N,T,D,H,W,C] 的 TensorShape。
     * 目前 seqLen 必须为 1。
     * @param s sequence length
     * @param r direction
     * @param n batch
     * @param t time
     * @param d depth
     * @param h height
     * @param w width
     * @param c channels
     */
    static create(s: number, r: number, n: number, t: number, d: number, h: number, w: number, c: number) {
        const shape = new TensorShape();
        shape.init(s, r, n, t, d, h, w, c);

        shape.usesNamedDimensions = NamedDimension.All;
        return shape;
    }

    init(s: number, r: number, n: number, t: number, d: number, h: number, w: number, c: number) {
        const data = this._data;
        data[SequenceLength] = s > 0 ? s : 1;
        data[NumberOfDirections] = r > 0 ? r : 1;
        data[DataBatch] = n > 0 ? n : 1;
        data[DataFeature3] = t > 0 ? t : 1;
        data[DataFeature2] = d > 0 ? d : 1;
        data[DataFeature1] = h > 0 ? h : 1;
        data[DataFeature0] = w > 0 ? w : 1;
        data[DataChannel] = c > 0 ? c : 1;
    }

    constructor() {
        this._data = new Array(MaxRank).fill(0);
    }

    //#region Helper functions

    /**
     *
     * @en
     * Allow to use negative axis to access tensorShape backward.
     *`axis` should be from -rank to rank (exclusive).
     *
     * @zh
     * 允许使用负轴访问张量形状的后面。
     * `axis` 应该从 -rank 到 rank（不包括）。
     * @param axis
     */
    axis(axis: number): number {
        assert(axis > -this.rank && axis < this.rank, `axis ${axis} is out of bounds`);

        return axis >= 0 ? axis : this.rank + axis;
    }

    /**
     * @en
     * Given an offset in memory return the dimensions indices of the element as [_,_,N,_,_,H,W,C].
     * @zh
     * 给定内存中的偏移量，返回元素的维度索引，如 [_,_,N,_,_,H,W,C]。
     * @param index
     */
    getPositionsFromIndex(index: number): [n: number, h: number, w: number, c: number] {
        let shape: TensorShape = null!;
        if (!this.hasNamedDimensions) {
            shape = this.asNamed();
        } else {
            // eslint-disable-next-line @typescript-eslint/no-this-alias
            shape = this;
        }

        const c = index % shape.channels;
        const w = (index / shape.channels) % shape.width;
        const h = (index / (shape.channels * shape.width)) % shape.height;
        const n =
            (index / (shape.channels * shape.width * shape.height * shape.depth * shape.extraDimension)) % shape.batch;
        return [n, h, w, c];
    }

    /**
     * @en
     * Given an offset in memory return the dimensions indices of the element as [S,R,N,T,D,H,W,C].
     * @zh
     * 给定内存中的偏移量，返回元素的维度索引，如 [S,R,N,T,D,H,W,C]。
     * @param index
     */
    getPositionsFromIndex8D(
        index: number
    ): [s: number, r: number, n: number, t: number, d: number, h: number, w: number, c: number] {
        let shape: TensorShape = null!;
        if (!this.hasNamedDimensions) {
            shape = this.asNamed();
        } else {
            // eslint-disable-next-line @typescript-eslint/no-this-alias
            shape = this;
        }

        const c = index % shape.channels;
        const w = (index / shape.channels) % shape.width;
        const h = (index / (shape.channels * shape.width)) % shape.height;
        const d = (index / (shape.channels * shape.width * shape.height)) % shape.depth;
        const t = (index / (shape.channels * shape.width * shape.height * shape.depth)) % shape.extraDimension;
        const n =
            (index / (shape.channels * shape.width * shape.height * shape.depth * shape.extraDimension)) % shape.batch;
        const r =
            (index / (shape.channels * shape.width * shape.height * shape.depth * shape.extraDimension * shape.batch)) %
            shape.numberOfDirections;
        const s =
            (index /
                (shape.channels *
                    shape.width *
                    shape.height *
                    shape.depth *
                    shape.extraDimension *
                    shape.batch *
                    shape.numberOfDirections)) %
            shape.sequenceLength;

        return [s, r, n, t, d, h, w, c];
    }

    /**
     * @en
     * Given an offset in memory return the dimensions indices of the element as [S,R,N,T,D,H,W,C] in ChannelFirst memory layout.
     * @zh
     * 给定内存中的偏移量，返回元素的维度索引，如 [S,R,N,T,D,H,W,C] 在 ChannelFirst 内存布局中。
     * @param index
     */
    getPositionsFromIndexChannelFirst8D(
        index: number
    ): [s: number, r: number, n: number, t: number, d: number, h: number, w: number, c: number] {
        let shape: TensorShape = null!;
        if (!this.hasNamedDimensions) {
            shape = this.asNamed();
        } else {
            // eslint-disable-next-line @typescript-eslint/no-this-alias
            shape = this;
        }

        const w = index % shape.width;
        const h = (index / shape.width) % shape.height;
        const d = (index / (shape.width * shape.height)) % shape.depth;
        const t = (index / (shape.width * shape.height * shape.depth)) % shape.extraDimension;
        const c = (index / (shape.width * shape.height * shape.depth * shape.extraDimension)) % shape.channels;
        const n =
            (index / (shape.width * shape.height * shape.depth * shape.extraDimension * shape.channels)) % shape.batch;
        const r =
            (index / (shape.width * shape.height * shape.depth * shape.extraDimension * shape.channels * shape.batch)) %
            shape.numberOfDirections;
        const s =
            (index /
                (shape.width *
                    shape.height *
                    shape.depth *
                    shape.extraDimension *
                    shape.channels *
                    shape.batch *
                    shape.numberOfDirections)) %
            shape.sequenceLength;
        return [s, r, n, t, d, h, w, c];
    }

    /**
     * @en
     *Given an offset in memory return the dimensions indices of the element as [_,_,N,_,_,H,W,C] in ChannelFirst format.
     * @zh
     * 给定内存中的偏移量，返回元素的维度索引，如 [N,H,W,C] 在 ChannelFirst 内存布局中。
     * @param index
     */
    getPositionsFromIndexChannelFirst(index: number): [n: number, h: number, w: number, c: number] {
        let shape: TensorShape = null!;
        if (!this.hasNamedDimensions) {
            shape = this.asNamed();
        } else {
            // eslint-disable-next-line @typescript-eslint/no-this-alias
            shape = this;
        }

        const w = index % shape.width;
        const h = (index / shape.width) % shape.height;
        const c = (index / (shape.width * shape.height * shape.depth * shape.extraDimension)) % shape.channels;
        const n =
            (index / (shape.width * shape.height * shape.depth * shape.extraDimension * shape.channels)) % shape.batch;
        return [n, h, w, c];
    }

    public indexWithBroadcast4D(n: number, h: number, w: number, c: number): number {
        let shape: TensorShape = null!;
        if (!this.hasNamedDimensions) {
            shape = this.asNamed();
        } else {
            // eslint-disable-next-line @typescript-eslint/no-this-alias
            shape = this;
        }
        n %= shape.batch;
        h %= shape.height;
        w %= shape.width;
        c %= shape.channels;
        return shape.index4D(n, h, w, c);
    }

    /**
     * @en
     * Given an element dimensions indices [S,R,N,T,D,H,W,C] with broadcast support, return this element offset in memory.
     * @zh
     * 给定元素维度索引 [S,R,N,T,D,H,W,C]，支持广播，返回此元素在内存中的偏移量。
     * @param s
     * @param r
     * @param n
     * @param t
     * @param d
     * @param h
     * @param w
     * @param c
     * @returns
     */
    public indexWithBroadcast8D(
        s: number,
        r: number,
        n: number,
        t: number,
        d: number,
        h: number,
        w: number,
        c: number
    ): number {
        let shape: TensorShape = null!;
        if (!this.hasNamedDimensions) {
            shape = this.asNamed();
        } else {
            // eslint-disable-next-line @typescript-eslint/no-this-alias
            shape = this;
        }
        s %= shape.sequenceLength;
        r %= shape.numberOfDirections;
        n %= shape.batch;
        t %= shape.extraDimension;
        d %= shape.depth;
        h %= shape.height;
        w %= shape.width;
        c %= shape.channels;
        return this.index8D(s, r, n, t, d, h, w, c);
    }

    public indexWithClamp4D(n: number, h: number, w: number, c: number): number {
        let shape: TensorShape = null!;
        if (!this.hasNamedDimensions) {
            shape = this.asNamed();
        } else {
            // eslint-disable-next-line @typescript-eslint/no-this-alias
            shape = this;
        }

        n = Math.max(n, 0);
        h = Math.max(h, 0);
        w = Math.max(w, 0);
        c = Math.max(c, 0);
        n = Math.min(n, shape.batch - 1);
        h = Math.min(h, shape.height - 1);
        w = Math.min(w, shape.width - 1);
        c = Math.min(c, shape.channels - 1);
        return this.index4D(n, h, w, c);
    }

    public indexWithClamp5D(n: number, d: number, h: number, w: number, c: number): number {
        let shape: TensorShape = null!;
        if (!this.hasNamedDimensions) {
            shape = this.asNamed();
        } else {
            // eslint-disable-next-line @typescript-eslint/no-this-alias
            shape = this;
        }

        n = Math.max(n, 0);
        d = Math.max(d, 0);
        h = Math.max(h, 0);
        w = Math.max(w, 0);
        c = Math.max(c, 0);
        n = Math.min(n, shape.batch - 1);
        d = Math.min(d, shape.depth - 1);
        h = Math.min(h, shape.height - 1);
        w = Math.min(w, shape.width - 1);
        c = Math.min(c, shape.channels - 1);
        return this.index5D(n, d, h, w, c);
    }

    /**
     * @en
     * Given an element dimensions indices [S,R,N,T,D,H,W,C] return this element offset in memory, clamping indices to tensor dimensions.
     * @zh
     * 给定元素维度索引 [S,R,N,T,D,H,W,C]，返回此元素在内存中的偏移量，将索引限制在张量维度内。
     * @param s
     * @param r
     * @param n
     * @param t
     * @param d
     * @param h
     * @param w
     * @param c
     * @returns
     */
    public indexWithClamp8D(
        s: number,
        r: number,
        n: number,
        t: number,
        d: number,
        h: number,
        w: number,
        c: number
    ): number {
        let shape: TensorShape = null!;
        if (!this.hasNamedDimensions) {
            shape = this.asNamed();
        } else {
            // eslint-disable-next-line @typescript-eslint/no-this-alias
            shape = this;
        }

        s = Math.max(s, 0);
        r = Math.max(r, 0);
        n = Math.max(n, 0);
        t = Math.max(t, 0);
        d = Math.max(d, 0);
        h = Math.max(h, 0);
        w = Math.max(w, 0);
        c = Math.max(c, 0);
        s = Math.min(s, shape.sequenceLength - 1);
        r = Math.min(r, shape.numberOfDirections - 1);
        n = Math.min(n, shape.batch - 1);
        t = Math.min(t, shape.extraDimension - 1);
        d = Math.min(d, shape.depth - 1);
        h = Math.min(h, shape.height - 1);
        w = Math.min(w, shape.width - 1);
        c = Math.min(c, shape.channels - 1);
        return this.index8D(s, r, n, t, d, h, w, c);
    }

    /**
     * @en
     * Given an element dimensions indices [S,R,N,T,D,H,W,C] return this element offset in memory.
     * @zh
     *
     * 给定元素维度索引 [S,R,N,T,D,H,W,C]，返回此元素在内存中的偏移量。
     * @param s
     * @param r
     * @param n
     * @param t
     * @param d
     * @param h
     * @param w
     * @param c
     * @returns
     */
    index8D(s: number, r: number, n: number, t: number, d: number, h: number, w: number, c: number): number {
        let shape: TensorShape = null!;
        if (!this.hasNamedDimensions) {
            shape = this.asNamed();
        } else {
            // eslint-disable-next-line @typescript-eslint/no-this-alias
            shape = this;
        }

        const index =
            c +
            w * shape.channels +
            h * shape.width * shape.channels +
            d * shape.height * shape.width * shape.channels +
            t * shape.depth * shape.height * shape.width * shape.channels +
            n * shape.extraDimension * shape.depth * shape.height * shape.width * shape.channels +
            r * shape.batch * shape.extraDimension * shape.depth * shape.height * shape.width * shape.channels +
            s *
            shape.numberOfDirections *
            shape.batch *
            shape.extraDimension *
            shape.depth *
            shape.height *
            shape.width *
            shape.channels;

        return index;
    }

    /**
     * @en
     * Given an element dimensions indices [0,0,N,0,D,H,W,C] return this element offset in memory.
     * @zh
     * 给定元素维度索引 [0,0,N,0,D,H,W,C]，返回此元素在内存中的偏移量。
     * @param n
     * @param d
     * @param h
     * @param w
     * @param c
     * @returns
     */
    index5D(n: number, d: number, h: number, w: number, c: number): number {
        let shape: TensorShape = null!;
        if (!this.hasNamedDimensions) {
            shape = this.asNamed();
        } else {
            // eslint-disable-next-line @typescript-eslint/no-this-alias
            shape = this;
        }

        const index =
            c +
            w * shape.channels +
            h * shape.width * shape.channels +
            d * shape.height * shape.width * shape.channels +
            n * shape.extraDimension * shape.depth * shape.height * shape.width * shape.channels;

        return index;
    }

    /**
     * @en
     * Given an element dimensions indices [0,0,N,0,0,H,W,C] return this element offset in memory.
     * @zh
     * 给定元素维度索引 [0,0,N,0,0,H,W,C]，返回此元素在内存中的偏移量。
     * @param n
     * @param h
     * @param w
     * @param c
     * @returns
     */
    index4D(n: number, h: number, w: number, c: number): number {
        let shape: TensorShape = null!;
        if (!this.hasNamedDimensions) {
            shape = this.asNamed();
        } else {
            // eslint-disable-next-line @typescript-eslint/no-this-alias
            shape = this;
        }
        const index =
            c +
            w * shape.channels +
            h * shape.width * shape.channels +
            n * shape.extraDimension * shape.depth * shape.height * shape.width * shape.channels;

        return index;
    }

    /**
     * @en
     * Given an element dimensions indices [S,R,N,T,D,H,W,C] return this element offset in memory in ChannelFirst format.
     * @zh
     * 给定元素维度索引 [S,R,N,T,D,H,W,C]，返回此元素在内存中的偏移量，以 ChannelFirst 格式。
     * @param s
     * @param r
     * @param n
     * @param t
     * @param d
     * @param h
     * @param w
     * @param c
     * @returns
     */
    indexChannelFirst8D(
        s: number,
        r: number,
        n: number,
        t: number,
        d: number,
        h: number,
        w: number,
        c: number
    ): number {
        let shape: TensorShape = null!;
        if (!this.hasNamedDimensions) {
            shape = this.asNamed();
        } else {
            // eslint-disable-next-line @typescript-eslint/no-this-alias
            shape = this;
        }
        const index =
            s *
            shape.numberOfDirections *
            shape.batch *
            shape.extraDimension *
            shape.depth *
            shape.height *
            shape.width *
            shape.channels +
            r * shape.batch * shape.extraDimension * shape.depth * shape.height * shape.width * shape.channels +
            n * shape.extraDimension * shape.depth * shape.height * shape.width * shape.channels +
            c * shape.extraDimension * shape.depth * shape.height * shape.width +
            t * shape.depth * shape.height * shape.width * +d * shape.height * shape.width +
            h * shape.width +
            w;
        return index;
    }

    /**
     * @en
     * Given an element dimensions indices [0,0,N,0,0,H,W,C] return this element offset in memory in ChannelFirst format.
     * @zh
     * 给定元素维度索引 [0,0,N,0,0,H,W,C]，返回此元素在内存中的偏移量，以 ChannelFirst 格式。
     * @param n
     * @param h
     * @param w
     * @param c
     */
    indexChannelFirst(n: number, h: number, w: number, c: number): number {
        let shape: TensorShape = null!;
        if (!this.hasNamedDimensions) {
            shape = this.asNamed();
        } else {
            // eslint-disable-next-line @typescript-eslint/no-this-alias
            shape = this;
        }
        const index =
            w +
            h * shape.width +
            c * shape.extraDimension * shape.depth * shape.width * shape.height +
            n * shape.channels * shape.extraDimension * shape.depth * shape.width * shape.height;

        return index;
    }

    /**
     * @en
     * Given an element dimensions indices [0,0,N,0,0,0,0,C] return this element offset in memory.
     *
     * @zh
     * 给定元素维度索引 [0,0,N,0,0,0,0,C]，返回此元素在内存中的偏移量。
     * @param n
     * @param c
     * @returns
     */
    index2D(n: number, c: number): number {
        let shape: TensorShape = null!;
        if (!this.hasNamedDimensions) {
            shape = this.asNamed();
        } else {
            // eslint-disable-next-line @typescript-eslint/no-this-alias
            shape = this;
        }
        const index = n * shape.flatWidth + c;

        return index;
    }

    //#endregion

    getAxis(axis: number): number {
        if (axis >= this.rank) {
            return 0;
        }

        if (this.hasNamedDimensions) {
            switch (axis) {
                case 0:
                    return this.sequenceLength;
                case 1:
                    return this.numberOfDirections;
                case 2:
                    return this.batch;
                case 3:
                    return this.extraDimension;
                case 4:
                    return this.depth;
                case 5:
                    return this.height;
                case 6:
                    return this.width;
                case 7:
                    return this.channels;
            }
        }

        return this._data[axis];
    }

    setAxis(axis: number, value: number) {
        if (this.hasNamedDimensions) {
            axis = axis < 0 || axis > 7 ? 7 : axis;
        } else {
            axis = this.axis(axis);
        }

        if (axis >= this.rank) {
            return;
        }

        if (this.hasNamedDimensions) {
            this._data[axis] = value > 0 ? value : 1;
        } else {
            this._data[axis] = value;
        }
    }

    toArray(): number[] {
        const size = this.rank;
        if (size > 0) {
            return this._data.slice(0, size);
        } else {
            return [1];
        }
    }

    /**
     * @en
     * Remove single-dimensional entries from the shape.
     * @zh
     * 从形状中删除单维条目。
     */
    squeeze(): TensorShape {
        let shape: TensorShape = null!;
        if (!this.hasNamedDimensions) {
            shape = this.asNamed();
        } else {
            // eslint-disable-next-line @typescript-eslint/no-this-alias
            shape = this;
        }

        const dims = shape.toArray();
        const squeezed = TensorShape.create(1, 1, 1, 1, 1, 1, 1, 1);
        assert(dims.length == squeezed.rank, 'Squeezed shape length must be equal to rank');
        let index = squeezed.rank;
        dims.forEach((dim) => {
            index--;
            if (dim > 1) {
                squeezed.setAxis(index, dim);
            }
        });

        return squeezed;
    }

    /**
     * @en
     * Return a TensorShape of dimensions [S,R,N,1,1,1,1,T*D*H*W*C].
     * @zh
     * 返回一个维度为 [S,R,N,1,1,1,1,T*D*H*W*C] 的 TensorShape。
     */

    flatten(): TensorShape {
        let shape: TensorShape = null!;
        if (!this.hasNamedDimensions) {
            shape = this.asNamed();
        } else {
            // eslint-disable-next-line @typescript-eslint/no-this-alias
            shape = this;
        }

        return TensorShape.create(
            shape.sequenceLength,
            shape.numberOfDirections,
            shape.batch,
            1,
            1,
            1,
            1,
            shape.flatWidth
        );
    }

    static equals(a: TensorShape, b: TensorShape): boolean {
        if (a.rank != b.rank) {
            return false;
        }

        for (let i = 0; i < a.rank; i++) {
            if (a.getAxis(i) != b.getAxis(i)) {
                return false;
            }
        }

        return true;
    }
}

export enum TensorDataType {
    Float,
    Half
}

export class Tensor {


    private _name: string = '';
    public get name(): string {
        return this._name;
    }
    public set name(value: string) {
        this._name = value;
    }

    private _tensorAllocator: ITensorAllocator = null!;
    public get tensorAllocator(): ITensorAllocator {
        return this._tensorAllocator;
    }
    public set tensorAllocator(value: ITensorAllocator) {
        this._tensorAllocator = value;
    }

    private _shape: TensorShape = null!;
    public get shape(): TensorShape {
        return this._shape;
    }
    public set shape(value: TensorShape) {
        this._shape = value;
    }

    get dateType(): TensorDataType {
        return this._preferredDataType;
    }

    set dateType(value: TensorDataType) {
        this._preferredDataType = value;
    }

    private _tensorOnDevice: ITensorData | null = null;

    public get sequenceLength(): number {
        return this.shape.sequenceLength;
    }
    public get numberOfDirections(): number {
        return this.shape.numberOfDirections;
    }
    public get batch(): number {
        return this.shape.batch;
    }
    public get extraDimension(): number {
        return this.shape.extraDimension;
    }
    public get depth(): number {
        return this.shape.depth;
    }
    public get height(): number {
        return this.shape.height;
    }
    public get width(): number {
        return this.shape.width;
    }
    public get channels(): number {
        return this.shape.channels;
    }
    public get kernelSpatialDepth(): number {
        return this.shape.kernelSpatialDepth;
    }
    public get kernelWidth(): number {
        return this.shape.kernelWidth;
    }
    public get kernelHeight(): number {
        return this.shape.kernelHeight;
    }
    public get kernelDepth(): number {
        return this.shape.kernelDepth;
    }
    public get kernelCount(): number {
        return this.shape.kernelCount;
    }
    public get flatHeight(): number {
        return this.shape.flatHeight;
    }
    public get flatWidth(): number {
        return this.shape.flatWidth;
    }
    public get length(): number {
        return this.shape.length;
    }
    public get dimensions(): number {
        return this.shape.dimensions;
    }

    _preferredDataType: TensorDataType = TensorDataType.Float;

    private _cache: number[] = [];
    public get cache(): number[] {
        return this._cache;
    }
    public set cache(value: number[]) {
        this._cache = value;
    }



    private _cacheIsDirty = false;

    getCache(
        indexArr:
            | number
            | [b: number, ch: number]
            | [b: number, h: number, w: number, ch: number]
            | [b: number, d: number, h: number, w: number, ch: number]
            | [s: number, r: number, n: number, t: number, d: number, h: number, w: number, c: number]
    ): number {
        this.prepareCacheForAccess();
        let index: number = 0;

        if (indexArr instanceof Array) {
            if (indexArr.length == 2) {
                index = this._shape.index2D(indexArr[0], indexArr[1]);
            } else if (indexArr.length == 4) {
                index = this._shape.index4D(indexArr[0], indexArr[1], indexArr[2], indexArr[3]);
            } else if (indexArr.length == 5) {
                index = this._shape.index5D(indexArr[0], indexArr[1], indexArr[2], indexArr[3], indexArr[4]);
            } else if (indexArr.length == 8) {
                index = this._shape.index8D(
                    indexArr[0],
                    indexArr[1],
                    indexArr[2],
                    indexArr[3],
                    indexArr[4],
                    indexArr[5],
                    indexArr[6],
                    indexArr[7]
                );
            }
        } else {
            index = indexArr;
        }

        return this._cache[index];
    }
    setCache(
        indexArr:
            | number
            | [b: number, ch: number]
            | [b: number, h: number, w: number, ch: number]
            | [b: number, d: number, h: number, w: number, ch: number]
            | [s: number, r: number, n: number, t: number, d: number, h: number, w: number, c: number],
        value: number
    ) {
        this.prepareCacheForAccess();
        let index: number = 0;

        if (indexArr instanceof Array) {
            if (indexArr.length == 2) {
                index = this._shape.index2D(indexArr[0], indexArr[1]);
            } else if (indexArr.length == 4) {
                index = this._shape.index4D(indexArr[0], indexArr[1], indexArr[2], indexArr[3]);
            } else if (indexArr.length == 5) {
                index = this._shape.index5D(indexArr[0], indexArr[1], indexArr[2], indexArr[3], indexArr[4]);
            } else if (indexArr.length == 8) {
                index = this._shape.index8D(
                    indexArr[0],
                    indexArr[1],
                    indexArr[2],
                    indexArr[3],
                    indexArr[4],
                    indexArr[5],
                    indexArr[6],
                    indexArr[7]
                );
            }
        } else {
            index = indexArr;
        }

        this._cache[index] = value;
        this._cacheIsDirty = true;
    }

    prepareCacheForAccess(): boolean {
        if (!this._cache) {
            if (this._tensorOnDevice != null) {
                this._cache = this._tensorOnDevice.download(this._shape);
            } else {
                this._cache = new Array(this.length).fill(0);
            }
            this._cacheIsDirty = false;
        }
        return true;
    }
}
